import os
import pandas as pd
from pathlib import Path
import dolphindb as ddb
from datetime import datetime

class DataImporter:
    def __init__(self):
        self.source_path = Path("F:/资料/历史数据/大商所")
        self.session = None
        self.upserter = None
        
    def connect_db(self):
        """连接DolphinDB"""
        self.session = ddb.session()
        self.session.connect("localhost", 8848, "admin", "123456")
        # 使用 TableUpserter 替代 TableAppender
        self.upserter = ddb.TableUpserter(
            dbPath="dfs://futures", 
            tableName="daily_all", 
            ddbSession=self.session,
            keyColNames=["date", "code", "exchange_code"]  # 指定键值列，用于判断重复
        )
        
    def process_file(self, file_path):
        """处理单个文件"""
        # 根据文件扩展名选择不同的读取方式
        ext = file_path.suffix.lower()
        
        try:
            if ext == '.csv':
                # 尝试不同的编码格式和分隔符
                encodings = ['gbk', 'gb2312', 'gb18030', 'utf-8']
                separators = [',', '\t', ';', '|']  # 常见的分隔符
                df = None
                
                for encoding in encodings:
                    for sep in separators:
                        try:
                            df = pd.read_csv(file_path, encoding=encoding, sep=sep)
                            # 如果成功读取并且列数正确，则跳出循环
                            if len(df.columns) > 1:
                                print(f"成功使用编码 {encoding} 和分隔符 '{sep}' 读取文件")
                                break
                        except (UnicodeDecodeError, pd.errors.EmptyDataError, pd.errors.ParserError):
                            continue
                    if df is not None and len(df.columns) > 1:
                        break
                        
                if df is None:
                    print(f"无法读取文件 {file_path}，所有编码格式和分隔符均失败")
                    return None
                    
            elif ext == '.xlsx':
                df = pd.read_excel(file_path)
            else:
                print(f"不支持的文件格式: {file_path}")
                return None
                
            # 打印列名，帮助调试
            print("读取到的列名:", df.columns.tolist())
            
            # 标准化列名
            if '合约' in df.columns:
                # 重命名列
                rename_dict = {
                    '日期': 'date',
                    '合约': 'code',
                    '最高价': 'high',
                    '最低价': 'low',
                    '开盘价': 'open',
                    '收盘价': 'close',
                    '前收盘价': 'pre_close',
                    '成交量': 'volume',
                    '成交金额' if '成交金额' in df.columns else '成交额': 'amount'
                }
                
                df = df.rename(columns=rename_dict)
                
                # 处理日期格式
                df['date'] = pd.to_datetime(df['date'], format='%Y%m%d')
                short_code = file_path.name.split('.')[0].split('_')[0]
                symbol_name = file_path.name.split('.')[0].split('_')[1]
                
                # 提取合约名称和交易所信息
                df['name'] = df['code'].apply(lambda x: symbol_name+x.replace(short_code, ""))
                df['exchange_code'] = 'DCE'
                df['exchange_name'] = '大商所'

                # 选择需要的列
                columns = ['date', 'code', 'name', 'exchange_code', 'exchange_name', 
                          'high', 'low', 'open', 'close', 'pre_close', 'volume', 'amount']
                return df[columns]
            
            return None
            
        except Exception as e:
            print(f"处理文件失败 {file_path}: {str(e)}")
            return None
            
    def import_to_db(self, df):
        """使用TableUpserter导入数据，自动处理重复数据"""
        try:
            # 直接使用upserter的upsert方法
            self.upserter.upsert(df)
            print(f"成功处理 {len(df)} 条数据")
        except Exception as e:
            print(f"导入数据失败: {str(e)}")
            # 如果是连接断开，尝试重连
            try:
                print("尝试重新连接...")
                self.connect_db()
                self.upserter.upsert(df)
                print(f"重连后成功处理 {len(df)} 条数据")
            except Exception as e2:
                print(f"重试失败: {str(e2)}")
            
    def process_all_files(self):
        """处理所有文件"""
        try:
            self.connect_db()
            
            # 遍历所有年份目录
            for year_dir in self.source_path.iterdir():
                if not year_dir.is_dir():
                    continue
                    
                print(f"\n处理 {year_dir.name} 年数据...")
                
                # 处理该年份下的所有文件
                for file_path in year_dir.glob('*.*'):
                    if file_path.suffix.lower() not in ['.csv', '.xlsx']:
                        continue
                        
                    print(f"处理文件: {file_path.name}")
                    df = self.process_file(file_path)
                    if df is not None and not df.empty:
                        self.import_to_db(df)
                        
        except Exception as e:
            print(f"处理过程出错: {str(e)}")
            raise
        finally:
            if self.session:
                self.session.close()

def main():
    """主函数"""
    importer = DataImporter()
    importer.process_all_files()

if __name__ == "__main__":
    main() 